import bisect
import os
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import List

import pandas as pd
import plotly.graph_objects as go

from utils import get_code_market_data


@dataclass
class OrderType:
    name: str
    size: int

    def __repr__(self):
        return f"[Order: {self.name} {self.size:.2f}]"


@dataclass
class PositionType:
    name: str
    size: int
    mean_value: float

    def add_size(self, size, price):
        if size * self.size > 0:  # 同号
            self.mean_value = (
                self.mean_value * abs(self.size) + abs(size) * price
            ) / abs(self.size + size)
            self.size = self.size + size
        else:
            if abs(self.size) < abs(size):
                self.mean_value = price
            elif abs(self.size) == abs(size):
                self.mean_value = 0
            self.size += size

    def __repr__(self):
        return f"[Position: {self.name} {self.size:.2f} {self.mean_value:.2f}]"


class BaseStrategy:
    def set_backtest_model(self, model):
        self.backtest = model

    def apply(self):
        pass

    def after_order(self, order: OrderType):
        pass

    def after_all(self):
        pass


class BacktestModel:
    def __init__(self):
        self.datas = {}
        self.index = None  # type: pd.Index
        self.start_date = None
        self.end_date = None
        self.cur_idx = None
        self.positions = {}
        self.cash = 0
        self.strategy = None
        self.commission_ratio = None
        self.commission_min = None
        self.origin_cash = None
        self.slip = None
        self.slip_cent = None
        self.margin_ratio = 1

    def reset(self):
        self.cash = self.origin_cash
        self.positions = {}
        self.index = pd.Index(data=[])
        for name, df in self.datas.items():
            self.index = self.index.union(df.index)
        self.index = self.index.sort_values()
        self.cur_idx = bisect.bisect_left(self.index, self.start_date)

    def set_origin_cash(self, cash):
        self.origin_cash = cash
        self.cash = cash

    def set_margin_ratio(self, ratio: float):
        self.margin_ratio = ratio

    def add_data(self, df, name):
        self.datas[name] = df

    def set_backtest_date(self, start_date: str | int, end_date: str | int):
        self.start_date = start_date
        self.end_date = end_date

    def set_commission(self, ratio: float = None, min_value: float = None):
        self.commission_ratio = ratio
        self.commission_min = min_value

    def set_slip(self, slip: float = None, slip_cent: float = None):
        self.slip = slip
        self.slip_cent = slip_cent

    def get_data(self, name, col, idx, default=None):
        if self.cur_idx is None:
            self.reset()
        if (
            name in self.datas
            and col in self.datas[name]
            and self.index[self.cur_idx + idx] in self.datas[name][col]
        ):
            return self.datas[name][col][self.index[self.cur_idx + idx]]
        else:
            return default

    def get_positions(self):
        return self.positions.copy()

    def get_cash(self):
        return self.cash

    def get_stock_value(self):
        return sum(
            [
                self.get_data(name, "close", 0, 0) * abs(pos.size) * self.margin_ratio
                for name, pos in self.positions.items()
            ]
        )

    def get_cur_time(self):
        return self.index[self.cur_idx]

    def get_value(self):
        return self.get_stock_value() + self.cash

    def add_strategy(self, strategy: BaseStrategy):
        self.strategy = strategy
        self.strategy.set_backtest_model(self)

    def deal_order(self, order: OrderType):
        name = order.name
        size = order.size
        if order.size == 0:
            return

        if name not in self.datas:
            raise Exception(f"{name} not in datas")

        if name not in self.positions:
            self.positions[name] = PositionType(name, 0, 0)
        position = self.positions.get(name)  # type: PositionType

        if size * position.size < 0 < abs(size) - abs(position.size):  # 平仓且超过
            old_size = position.size
            self.deal_order(OrderType(order.name, -old_size))
            self.deal_order(OrderType(order.name, size + old_size))
            return

        cur_price = self.get_data(name, "close", 0)
        if not isinstance(cur_price, (int, float)):
            raise Exception(f"price {cur_price} is invalid")

        if size * position.size < 0:  # 平仓
            if self.slip_cent:
                price = (1 - self.slip_cent) * cur_price
            elif self.slip:
                price = cur_price - self.slip
            else:
                price = cur_price
        else:
            if self.slip_cent:
                price = (1 + self.slip_cent) * cur_price
            elif self.slip:
                price = cur_price + self.slip
            else:
                price = cur_price

        if self.commission_ratio:
            commission = max(
                self.commission_ratio * abs(size) * price,
                self.commission_min if self.commission_min else 0,
            )
        else:
            commission = self.commission_min if self.commission_min else 0

        if size * position.size < 0:  # 如果平仓
            change_cash = abs(size) * price * self.margin_ratio - commission
        else:  # 如果只是加仓
            change_cash = -abs(size) * price * self.margin_ratio - commission

        if self.cash + change_cash < 0:
            raise Exception(f"Not enough cash for order")
        self.cash += change_cash

        position.add_size(size, price)

        if abs(self.positions[name].size) < 0.01:
            self.positions.pop(name)

    def order(self, order: OrderType):
        self.deal_order(order)
        if self.strategy:
            self.strategy.after_order(order)

    def slide(self):
        self.cur_idx += 1
        return (
            self.cur_idx >= len(self.index) - 1
            or self.index[self.cur_idx] > self.end_date
        )

    def change_cash(self, value):
        self.cash += value

    def run_backtest(self):
        self.reset()
        while self.cur_idx < len(self.index):
            self.strategy.apply()
            self.cur_idx += 1
            if (
                self.cur_idx >= len(self.index)
                or self.index[self.cur_idx] > self.end_date
            ):
                break

        self.strategy.after_all()


class Strategy(BaseStrategy):
    def __init__(self, consider_stocks: List, stock_num: int, save_dir: str):
        self.consider_stocks = consider_stocks  # type: List[str]
        self.stock_num = stock_num
        self.save_dir = save_dir
        self.time_list = []
        self.value_list = []
        os.makedirs(self.save_dir, exist_ok=True)

    def after_all(self):
        fig = go.Figure()
        fig.add_trace(go.Scatter(x=self.time_list, y=self.value_list, mode="lines"))
        fig.update_layout(
            title="收益曲线图",
            xaxis_title="日期",
            yaxis_title="价值",
            xaxis_type="category",
            xaxis_rangeslider_visible=True,
        )
        fig.write_html(os.path.join(self.save_dir, "profit.html"))

        df = pd.DataFrame({"datetime": self.time_list, "value": self.value_list})
        df.to_csv(os.path.join(self.save_dir, "daily_value.csv"))

    def after_order(self, order):
        print(
            f"时间：{self.backtest.get_cur_time()}, {order} \t 资产:{self.backtest.get_value():.2f}, "
            f"标的:{self.backtest.get_stock_value():.2f}, 现金:{self.backtest.get_cash():.2f}"
            f"仓位: {self.backtest.positions}"
        )

    def apply(self):
        import random

        size = random.choice([-300, -600, 200, 400, 300])
        self.backtest.order(OrderType("603909.SH", size))

    def apply2(self):
        for name in self.consider_stocks:
            if self.backtest.get_data(name, "suspendFlag", 0, 0):
                print(f"suspend: {self.backtest.get_cur_time()} {name}")
            if self.backtest.get_data(name, "suspendFlag", 0, 0) is None:
                print(f"None suspend")
        today_close = {
            name: self.backtest.get_data(name, "close", 0)
            for name in self.consider_stocks
            if self.backtest.get_data(name, "close", 0) is not None
        }
        today_close_pre_10 = {
            name: self.backtest.get_data(name, "close", -10)
            for name in self.consider_stocks
            if self.backtest.get_data(name, "close", -10) is not None
        }
        use_stocks = set(today_close.keys()).intersection(
            set(today_close_pre_10.keys())
        )
        today_change = {
            name: (today_close.get(name) - today_close_pre_10.get(name))
            / today_close_pre_10.get(name)
            for name in use_stocks
        }
        sort_name_ratio = sorted(
            today_change.items(), key=lambda x: x[1], reverse=False
        )
        sort_names = [x[0] for x in sort_name_ratio[: self.stock_num]]
        cur_value = self.backtest.get_value()
        one_stock_value = cur_value / self.stock_num

        cur_positions = self.backtest.get_positions()
        for name, x in cur_positions.items():
            self.backtest.order(OrderType(name, -x.size))
        target_positions = {
            name: one_stock_value // (today_close.get(name) * 1) * 1
            for name in sort_names
        }
        for name, size in target_positions.items():
            self.backtest.order(OrderType(name, size))
        self.time_list.append(self.backtest.get_cur_time())
        self.value_list.append(self.backtest.get_value())


def job(start_date, end_date, dividend_type):
    print("start job:", start_date, end_date, dividend_type)
    dividend_type = dividend_type
    # stocks = [
    #     "600001.SH",
    #     "600002.SH",
    #     "600003.SH",
    #     "600004.SH" "600005.SH",
    #     "600006.SH" "600007.SH",
    #     "600008.SH",
    #     "600009.SH",
    #     "600010.SH",
    # ]
    from xtquant.xtdata import get_stock_list_in_sector

    stocks = get_stock_list_in_sector("上证A股")[:50]
    stock_num = 5
    cash = 1e7
    save_dir = f"back_data/{start_date}_{end_date}/{dividend_type}"
    data_start = (
        datetime.strptime(start_date, "%Y%m%d") - timedelta(days=30)
    ).strftime("%Y%m%d")
    end_date = (datetime.strptime(end_date, "%Y%m%d") + timedelta(days=1)).strftime(
        "%Y%m%d"
    )
    back_model = BacktestModel()
    back_model.set_slip(slip_cent=0.001)
    back_model.set_margin_ratio(0.5)
    back_model.set_commission(0.001, 5)
    back_model.set_backtest_date(start_date, end_date)
    for stock in stocks:
        df = get_code_market_data(
            stock, data_start, end_date, period="1d", dividend_type=dividend_type
        )
        if len(df) > 0:
            back_model.add_data(df, stock)
    strategy = Strategy(stocks, stock_num, save_dir)
    back_model.add_strategy(strategy)
    back_model.set_origin_cash(cash)
    back_model.run_backtest()


if __name__ == "__main__":
    # "none" "front" "back" "front_ratio" "back_ratio"
    job("20201231", "20211231", "front_ratio")
